import hashlib
import hmac
import datetime
from typing import Dict

class JimengSignatureHelper:
    """即梦AI签名工具类"""
    
    def __init__(self, access_key: str, secret_key: str, host: str, region: str, service: str):
        self.access_key = access_key
        self.secret_key = secret_key
        self.host = host
        self.region = region
        self.service = service
    
    def _sign(self, key: bytes, msg: str) -> bytes:
        """HMAC-SHA256签名"""
        return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()
    
    def _get_signature_key(self, secret_key: str, date_stamp: str, region_name: str, service_name: str) -> bytes:
        """获取签名密钥"""
        k_date = self._sign(secret_key.encode('utf-8'), date_stamp)
        k_region = self._sign(k_date, region_name)
        k_service = self._sign(k_region, service_name)
        k_signing = self._sign(k_service, 'request')
        return k_signing
    
    def _format_query(self, parameters: Dict[str, str]) -> str:
        """格式化查询参数 - 与官方代码保持一致，不进行URL编码"""
        request_parameters_init = ''
        for key in sorted(parameters):
            request_parameters_init += key + '=' + parameters[key] + '&'
        request_parameters = request_parameters_init[:-1]
        return request_parameters
    
    def generate_v4_signature(self, method: str, query_params: Dict[str, str], body: str) -> Dict[str, str]:
        """生成火山引擎V4签名"""
        if not self.access_key or not self.secret_key:
            raise Exception("Missing access_key or secret_key for Jimeng AI")
        
        # 时间戳
        t = datetime.datetime.utcnow()
        current_date = t.strftime('%Y%m%dT%H%M%SZ')
        date_stamp = t.strftime('%Y%m%d')
        
        # 请求信息
        canonical_uri = '/'
        canonical_querystring = self._format_query(query_params)
        signed_headers = 'content-type;host;x-content-sha256;x-date'
        payload_hash = hashlib.sha256(body.encode('utf-8')).hexdigest()
        content_type = 'application/json'
        
        # 标准请求
        canonical_headers = (
            f'content-type:{content_type}\n'
            f'host:{self.host}\n'
            f'x-content-sha256:{payload_hash}\n'
            f'x-date:{current_date}\n'
        )
        
        canonical_request = (
            f'{method}\n{canonical_uri}\n{canonical_querystring}\n'
            f'{canonical_headers}\n{signed_headers}\n{payload_hash}'
        )
        
        # 签名字符串
        algorithm = 'HMAC-SHA256'
        credential_scope = f'{date_stamp}/{self.region}/{self.service}/request'
        string_to_sign = (
            f'{algorithm}\n{current_date}\n{credential_scope}\n'
            f'{hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()}'
        )
        
        # 计算签名
        signing_key = self._get_signature_key(self.secret_key, date_stamp, self.region, self.service)
        signature = hmac.new(signing_key, string_to_sign.encode('utf-8'), hashlib.sha256).hexdigest()
        
        # 授权头
        authorization_header = (
            f'{algorithm} Credential={self.access_key}/{credential_scope}, '
            f'SignedHeaders={signed_headers}, Signature={signature}'
        )
        
        return {
            'X-Date': current_date,
            'Authorization': authorization_header,
            'X-Content-Sha256': payload_hash,
            'Content-Type': content_type
        }